package step

import (
	"bytes"
	"crypto/sha256"
	"errors"
	"fmt"
	"github.com/wanchain/go-mpc/common/hexutil"
	"github.com/wanchain/go-mpc/log"
	"github.com/wanchain/go-mpc/p2p/discover"
	"github.com/wanchain/go-mpc/storeman/osmconf"
	schcomm "github.com/wanchain/go-mpc/storeman/schnorrcomm"
	mpcprotocol "github.com/wanchain/go-mpc/storeman/storemanmpc/protocol"
	"math/big"
	"strconv"
)

type MpcRRcvJudgeStepUni struct {
	BaseStep
	rcvColInterMap map[*discover.NodeID]*big.Int
	rcvColInter    *big.Int
	rcvCol         *big.Int
	sigNum         int
}

func CreateMpcRRcvJudgeStepUni(totalNum int, sigNum int) *MpcRRcvJudgeStepUni {
	log.SyslogInfo("CreateMpcRRcvJudgeStep begin")

	mpc := &MpcRRcvJudgeStepUni{
		*CreateBaseStep(totalNum, -1),
		make(map[*discover.NodeID]*big.Int, 0),
		nil, nil, sigNum}
	return mpc
}

func (ptStep *MpcRRcvJudgeStepUni) InitStep(result mpcprotocol.MpcResultInterface) error {
	log.SyslogDebug("MpcRRcvJudgeStepUni::InitStep begin", "ctx", ptStep.GetMpcContextId())
	ptStep.BaseStep.InitStep(result)

	ret, err := result.GetValue(mpcprotocol.RRcvedCollInter)
	if err != nil {
		log.SyslogErr("MpcRRcvJudgeStepUni", "InitStep.getValue error", err.Error(), "ctx", ptStep.GetMpcContextId())
		return err
	}
	if len(ret) == 0 {
		log.SyslogErr("MpcRRcvJudgeStepUni", "GetValue len(ret)", len(ret), "ctx", ptStep.GetMpcContextId())
		return err
	}
	//ptStep.rcvColInter = &ret[0]
	ptStep.rcvColInter = big.NewInt(0).SetBytes(ret[0].Bytes())

	ret, err = result.GetValue(mpcprotocol.RRcvedColl)
	if err != nil {
		log.SyslogErr("MpcRRcvJudgeStepUni", "InitStep.getValue error", err.Error(), "ctx", ptStep.GetMpcContextId())
		return err
	}
	if len(ret) == 0 {
		log.SyslogErr("MpcRRcvJudgeStepUni", "GetValue len(ret)", len(ret), "ctx", ptStep.GetMpcContextId())
		return err
	}
	//ptStep.rcvCol = &ret[0]
	ptStep.rcvCol = big.NewInt(0).SetBytes(ret[0].Bytes())

	log.SyslogDebug("MpcRRcvJudgeStepUni::InitStep end", "RRcvedColl", hexutil.Encode(ptStep.rcvCol.Bytes()), "ctx", ptStep.GetMpcContextId())
	log.SyslogDebug("MpcRRcvJudgeStepUni::InitStep end", "RRcvedCollInter", hexutil.Encode(ptStep.rcvColInter.Bytes()), "ctx", ptStep.GetMpcContextId())
	log.SyslogDebug("MpcRRcvJudgeStepUni::InitStep end", "ctx", ptStep.GetMpcContextId())
	return nil
}

func (ptStep *MpcRRcvJudgeStepUni) CreateMessage() []mpcprotocol.StepMessage {
	log.SyslogDebug("MpcRRcvJudgeStepUni::CreateMessage begin", "ctx", ptStep.GetMpcContextId())
	message := make([]mpcprotocol.StepMessage, 1)
	message[0].MsgCode = mpcprotocol.MPCMessage
	message[0].PeerID = nil

	var buf bytes.Buffer
	buf.Write(ptStep.rcvColInter.Bytes())
	h := sha256.Sum256(buf.Bytes())

	prv, err := osmconf.GetOsmConf().GetSelfPrvKey()
	if err != nil {
		return nil
	}
	r, s, _ := schcomm.SignInternalData(prv, h[:])

	message[0].Data = make([]big.Int, 3)
	message[0].Data[0] = *ptStep.rcvColInter
	message[0].Data[1] = *r
	message[0].Data[2] = *s
	log.SyslogDebug("MpcRRcvJudgeStepUni::CreateMessage end", "ctx", ptStep.GetMpcContextId())
	return message
}

func (ptStep *MpcRRcvJudgeStepUni) HandleMessage(msg *mpcprotocol.StepMessage) bool {
	log.SyslogInfo("MpcRRcvJudgeStepUni::HandleMessage begin", "ctx", ptStep.GetMpcContextId())
	r := msg.Data[1]
	s := msg.Data[2]

	var buf bytes.Buffer
	buf.Write(msg.Data[0].Bytes())
	h := sha256.Sum256(buf.Bytes())

	_, grpIdStr, err := ptStep.BaseStep.GetGrpId()
	if err != nil {
		log.SyslogErr("MpcRRcvJudgeStepUni", "HandleMessage error", err.Error(), "ctx", ptStep.GetMpcContextId())
	}

	senderPk, err := osmconf.GetOsmConf().GetPKByNodeId(grpIdStr, msg.PeerID)
	if err != nil {
		log.SyslogErr("MpcRRcvJudgeStepUni", "GetPKByNodeId error", err.Error(), "ctx", ptStep.GetMpcContextId())
	}

	bVerifySig := schcomm.VerifyInternalData(senderPk, h[:], &r, &s)

	if bVerifySig {
		log.SyslogInfo("MpcRRcvJudgeStepUni::HandleMessage check sig success", "ctx", ptStep.GetMpcContextId())
		ptStep.rcvColInterMap[msg.PeerID] = &msg.Data[0]
	} else {
		log.SyslogErr("MpcRRcvJudgeStepUni::HandleMessage check sig fail", "ctx", ptStep.GetMpcContextId())
	}

	log.SyslogInfo("MpcRRcvJudgeStepUni::HandleMessage Intersection collection information",
		" self Intersection collection", hexutil.Encode(ptStep.rcvColInter.Bytes()),
		"peerID", msg.PeerID.SlimString(),
		"peers intersection collection", hexutil.Encode(msg.Data[0].Bytes()), "ctx", ptStep.GetMpcContextId())

	log.SyslogInfo("MpcRRcvJudgeStepUni::HandleMessage end", "ctx", ptStep.GetMpcContextId())

	return true
}

func (ptStep *MpcRRcvJudgeStepUni) FinishStep(result mpcprotocol.MpcResultInterface, mpc mpcprotocol.StoremanManager) error {

	sigNum := ptStep.sigNum

	log.SyslogDebug("MpcRRcvJudgeStepUni::FinishStep begin", "ctx", ptStep.GetMpcContextId())

	// compute the intersec and save
	err := ptStep.BaseStep.FinishStep()
	if err != nil {
		return err
	}

	_, grpIdString, err := ptStep.BaseStep.GetGrpId()
	if err != nil {
		log.SyslogErr("MpcRRcvJudgeStepUni", "error in GetGrpId", err.Error(), "ctx", ptStep.GetMpcContextId())
		return err
	}

	leaderIndex, _ := osmconf.GetOsmConf().GetLeaderIndex(grpIdString)
	bIncludeLeader, err := osmconf.IsHaming(ptStep.rcvColInter, leaderIndex)
	if err != nil {
		log.SyslogErr("MpcRRcvJudgeStepUni", "error in IsHaming", err.Error(), "ctx", ptStep.GetMpcContextId())
		return err
	}
	// 1. should include leaderIndex
	if !bIncludeLeader {
		log.SyslogErr("MpcRRcvJudgeStepUni leader is not included in the intersection", "ctx", ptStep.GetMpcContextId())
		return errors.New(fmt.Sprintf("leader is not included in the intersection"))
	} else {
		log.SyslogInfo("self  rcvColInter include leader index", "ctx", ptStep.GetMpcContextId())
	}

	// 2. others' inter collection should be equal to self's inter collection
	for nodeId, rcvColInter := range ptStep.rcvColInterMap {
		if ptStep.rcvColInter.Cmp(rcvColInter) != 0 {
			log.SyslogErr("MpcRRcvJudgeStepUni received colInter not equal to self ",
				"received", hexutil.Encode(rcvColInter.Bytes()),
				"self", hexutil.Encode(ptStep.rcvColInter.Bytes()), "drop nodeId", nodeId.SlimString(), "ctx", ptStep.GetMpcContextId())
			//if not equal, drop the data from the peer.
			ptStep.MpcContext.InsertDropPeers([]discover.NodeID{*nodeId})
		}
	}
	// drop the peer not in the collections
	nodesNeedDrop, waitCount, err := ptStep.BaseStep.GetNodeIdsNotInColletion(grpIdString, ptStep.rcvColInter)
	if err != nil {
		return err
	}

	ptStep.MpcContext.InsertDropPeers(nodesNeedDrop)
	log.SyslogInfo("MpcRRcvJudgeStepUni::FinishStep",
		"waitCount", waitCount,
		"InsertDropPeers drop nodeId", discover.ArrNodeIdsToStr(nodesNeedDrop),
		"self rcv collection intersection", hexutil.Encode(ptStep.rcvColInter.Bytes()),
		"ctx", ptStep.GetMpcContextId())

	if waitCount > 0 {
		ptStep.MpcContext.UptAllStepsWtCount(waitCount)
	}

	// 3. RRcvedColl == RRcvedCollInter
	log.SyslogInfo("MpcRRcvJudgeStepUni::FinishStep", "rrcvedCol", hexutil.Encode(ptStep.rcvCol.Bytes()), "rcvColInter", hexutil.Encode(ptStep.rcvColInter.Bytes()), "ctx", ptStep.GetMpcContextId())
	if ptStep.rcvCol.Cmp(ptStep.rcvColInter) == 0 {
		log.SyslogInfo("MpcRRcvJudgeStepUni rrcvedCol is equal rcvColInter", "rrcvedCol", hexutil.Encode(ptStep.rcvCol.Bytes()), "rcvColInter", hexutil.Encode(ptStep.rcvColInter.Bytes()), "ctx", ptStep.GetMpcContextId())
		//return nil
	} else {
		log.SyslogInfo("MpcRRcvJudgeStepUni::FinishStep rrcvedCol is NOT equal rcvColInter ", "rrcvedCol", hexutil.Encode(ptStep.rcvCol.Bytes()), "rcvColInter", hexutil.Encode(ptStep.rcvColInter.Bytes()), "ctx", ptStep.GetMpcContextId())
		//return errors.New(fmt.Sprintf("rrcvedCol is NOT equal rcvColInter"))
	}

	// 4. RRcvedColl != RRcvedCollInter
	// 5. build rskShare and rpkShare by new intersection.
	totalNumber, err := osmconf.GetOsmConf().GetTotalNum(grpIdString)
	if err != nil {
		log.SyslogErr("MpcRRcvJudgeStepUni::FinishStep", "getTotalNumber err", err.Error(), "ctx", ptStep.GetMpcContextId())
		return err
	}

	bigs := make([]big.Int, 0)
	bigsIndexRpkShare := make([]big.Int, 0)
	for i := 0; i < int(totalNumber); i++ {
		b, err := osmconf.IsHaming(ptStep.rcvColInter, uint16(i))
		if err != nil {
			log.SyslogErr("MpcRRcvJudgeStepUni::FinishStep", "IsHaming err", err.Error(), "ctx", ptStep.GetMpcContextId())
			return err
		}
		if b {

			for sigIndex := 0; sigIndex < sigNum; sigIndex++ {
				key := mpcprotocol.RSKSIJ + strconv.Itoa(int(i)) + strconv.Itoa(sigIndex)
				sij, err := result.GetValue(key)
				if err != nil {
					log.SyslogErr("MpcRRcvJudgeStepUni::FinishStep", "get sij err", err.Error(), "key", key, "ctx", ptStep.GetMpcContextId())
					return err
				}
				bigs = append(bigs, sij[0])
			}
			bigsIndexRpkShare = append(bigsIndexRpkShare, *big.NewInt(int64(i)))
		}
	}

	for sigIndex := 0; sigIndex < sigNum; sigIndex++ {
		// get rskShare
		rskShare := big.NewInt(0)
		for _, value := range bigs[sigIndex*len(bigsIndexRpkShare) : (sigIndex+1)*len(bigsIndexRpkShare)] {
			rskShare.Add(rskShare, &value)
			rskShare.Mod(rskShare, ptStep.schnorrMpcer.GetMod())
		}

		key := mpcprotocol.RSkShare + strconv.Itoa(sigIndex)
		err = result.SetValue(key, []big.Int{*rskShare})
		if err != nil {
			return err
		}
		scmper := ptStep.schnorrMpcer

		// add for schnorrr-btc 340 begin
		if int(big.NewInt(0).SetBytes(ptStep.GetMpcContext().GetInput().Alg).Uint64()) == mpcprotocol.ALGSCHNORR340 {
			key = mpcprotocol.RSkShareInv + strconv.Itoa(sigIndex)
			rskShareInv, err := scmper.NSubX(rskShare)
			if err != nil {
				return err
			}
			err = result.SetValue(key, []big.Int{*rskShareInv})
			if err != nil {
				return err
			}

		}
		// add for schnorrr-btc 340 end

		rpkShare, err := scmper.SkG(rskShare)
		if err != nil {
			return err
		}

		selfIndex, err := osmconf.GetOsmConf().GetSelfInx(grpIdString)
		if err != nil {
			log.SyslogErr("MpcRRcvJudgeStepUni::FinishStep", "FinishStep", err.Error(), "ctx", ptStep.GetMpcContextId())
			return err
		}
		key = mpcprotocol.RPkShare + strconv.Itoa(int(selfIndex)) + strconv.Itoa(sigIndex)

		rpkShareBytes, err := scmper.MarshPt(rpkShare)
		if err != nil {
			return err
		}
		err = result.SetByteValue(key, rpkShareBytes)
		if err != nil {
			return err
		}

		colleciontRpkShare, _ := osmconf.BuildDataByIndexes(&bigsIndexRpkShare)
		log.SyslogInfo("MpcRRcvJudgeStepUni::FinishStep save(update) rpkShare",
			"key", key,
			"rpkShare", hexutil.Encode(rpkShareBytes),
			"selfInterCollection", hexutil.Encode(ptStep.rcvColInter.Bytes()),
			"sijCollection", hexutil.Encode(colleciontRpkShare.Bytes()),
			"ctx", ptStep.GetMpcContextId())

		log.SyslogInfo("MpcRRcvJudgeStepUni::FinishStep end", "ctx", ptStep.GetMpcContextId())
	}
	return nil
}
